# __init__.py
import os
import sys
from pathlib import Path

# 使用传统方式，但更加健壮
PROJECT_ROOT_STR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT_STR not in sys.path:
    sys.path.insert(0, PROJECT_ROOT_STR)

"""PPQ量化配置管理模块"""
import os
from dataclasses import dataclass
from typing import Dict, List
from ppq import TargetPlatform
from .path_config import DATA_DIR,MODELS_DIR

@dataclass
class QuantizationConfig:
    """量化配置参数类"""
    device: str = 'cuda'
    quant_platform: TargetPlatform = TargetPlatform.TRT_INT8
    input_shapes: Dict[str, List[int]] = None
    batch_size: int = 1
    calib_steps: int = 32
    num_workers: int = 8
    
    # 文件路径配置
    onnx_path: str = ""
    onnx_output_path: str = ""
    ppq_json_output_path: str = ""
    data_dir: str = ""
    
    # 量化设置
    fusion: bool = False
    lsq_optimization: bool = False
    
    def __post_init__(self):
        
        if self.input_shapes is None:
            self.input_shapes = {'images': [1, 3, 640, 640]}
        
        if not self.onnx_path:
            self.onnx_path = os.path.join(MODELS_DIR, 'onnx/yolov11n.onnx')
        if not self.onnx_output_path:
            self.onnx_output_path = os.path.join(MODELS_DIR, 'onnx/weights/yolov11n_int8.onnx')
        if not self.ppq_json_output_path:
            self.ppq_json_output_path = os.path.join(MODELS_DIR, 'onnx/weights/yolov11n_int8.json')
        if not self.data_dir:
            self.data_dir = os.path.join(DATA_DIR, 'coco128/images/train2017')

def create_default_config() -> QuantizationConfig:
    """创建默认配置"""
    return QuantizationConfig()